# coding:utf-8
# __user__ = hiicy redldw
# __time__ = 2019/8/15
# __file__ = augment
# __desc__ =
import imgaug.augmenters as iaa
import imgaug as ia
import os
import random
import numpy as np
from scipy import misc

"""
# 數據增強
datagen = ImageDataGenerator(
    rotation_range=20,
    # shear_range=0.1,  # 剪切强度（逆时针方向的剪切变换角度）
    fill_mode="nearest",  # ‘constant’‘nearest’，‘reflect’或‘wrap’之一，当进行变换时超出边界的点将根据本参数给定的方法进行处理
    cval=0,  # 浮点数或整数，当fill_mode=constant时，指定要向超出边界的点填充的值
    width_shift_range=0.2,  # 图片宽度的某个比例 图片水平偏移的幅度
    height_shift_range=0.2,  # 图片竖直偏移的幅度
    horizontal_flip=True,  # 进行随机水平翻转
    vertical_flip=True,  # 随机竖直翻转
    zoom_range=0.2,  # 浮点数或形如[lower,upper]的列表，随机缩放的幅度，若为浮点数，则相当于[lower,upper] = [1 - zoom_range, 1+zoom_range]
    channel_shift_range=30,
    rescale=1 / 255  # 重放缩因子
)


# 加噪音
def add_noise(oridir, tardir):
    for i in os.listdir(oridir):
        imgpath = os.path.join(oridir, i)
        img = misc.imread(imgpath)
        img = cv2.resize(img, (COLS, ROWS))
        noise = su.random_noise(img, mode='speckle', seed=2, clip=True, mean=0.1)
        misc.imsave(os.path.join(tardir, i[:-4] + '_speckle.png'), noise)


def do():
    from multiprocessing import Process
    dirss = [original_dataset_ONE_dir,
             original_dataset_TWO_dir,
             original_dataset_THREE_dir,
             original_dataset_OTHER_dir, ]
    ssdir = [Augument_dataset_ONE_dir,
             Augument_dataset_TWO_dir,
             Augument_dataset_THREE_dir,
             Augument_dataset_OTHER_dir]
    ps = []
    for sdir, tdir in zip(dirss[-1:], ssdir[-1:]):
        p = Process(target=add_noise, args=(sdir, tdir))
        ps.append(p)
    for p in ps:
        print("启动进程")
        p.start()
    for p in ps:
        p.join()


# do()
def delete():
    ssdir = [Augument_dataset_ONE_dir,
             Augument_dataset_TWO_dir,
             Augument_dataset_THREE_dir,
             ]
    for d in ssdir:
        for i in os.listdir(d):
            if 'noise' in i:
                imgpath = os.path.join(d, i)
                os.remove(imgpath)


# delete()
# do()
"""

class IAU:
    def __init__(self,ddir,output):
        self.dir = ddir
        self.output = output
    @property
    def imgs(self):
        imgs = []
        for iname in os.listdir(self.dir):
            imgs.append(os.path.join(self.dir,iname))
        return np.array(imgs)
    def _len(self):
        return len(self.imgs)
    def _random_name(self,prefix="ocr"):
        names = "1234567890sfafaezgthnct"
        alpha = "".join([chr(i) for i in range(97,123)])
        bigalpha = "".join([chr(i) for i in range(65,91)])

        _na = "".join(names[random.randint(0,10)] for _ in range(4))
        _al = "".join(alpha[random.randint(0,10)] for _ in range(4))
        _bl = "".join(bigalpha[random.randint(0,10)] for _ in range(4))

        return os.path.join(self.output,_al+"_"+_na+"_"+_bl)

    def _load_batch(self,batch=16):
        cbatch = [random.randint(0,self._len()-1) for _ in range(batch)]
        nimages = self.imgs[cbatch]
        v = []
        for image in nimages:
            img = misc.imread(image)
            v.append(img)
        return np.array(v)
    @property
    def seq(self):
        sometimes = lambda aug:iaa.Sometimes(0.7,aug)
        return iaa.Sequential(
            [
                sometimes(iaa.CropAndPad(
                    percent=(-0.05, 0.1),
                    pad_mode=ia.ALL,
                    pad_cval=(0, 255)
                )),
                sometimes(iaa.Affine(
                scale={"x": (0.8, 1.2), "y": (0.8, 1.2)},
                rotate=(-45, 45),  # rotate by -45 to +45 degrees
                shear=(-16,16),
                order=[0,1],
                )),
                iaa.SomeOf((0,5),
                           [
                               iaa.OneOf([
                                   iaa.GaussianBlur((0,3.0)),
                                   iaa.AverageBlur(k=(2,7)),
                                   iaa.MedianBlur(k=(3,11)),
                               ]),
                               iaa.Sharpen(alpha=(0,1.),lightness=(0.75,1.25)),
                               iaa.SimplexNoiseAlpha(iaa.OneOf([
                                   iaa.EdgeDetect(alpha=(0.5,1.0)),
                               ])),
                               # add gaussian noise to images
                               iaa.AdditiveGaussianNoise(loc=0, scale=(0.0, 0.05 * 255), per_channel=0.5),
                               iaa.Add((-10, 10), per_channel=0.5), # change brightness of images (by -10 to 10 of original value)
                               iaa.AddToHueAndSaturation((-20, 20)),  # change hue and saturation
                               # either change the brightness of the whole image (sometimes
                               # per channel) or change the brightness of subareas
                               iaa.OneOf([
                                   iaa.Multiply((0.5, 1.5), per_channel=0.5),
                                   iaa.FrequencyNoiseAlpha(
                                       exponent=(-4, 0),
                                       first=iaa.Multiply((0.5, 1.5), per_channel=True),
                                       second=iaa.ContrastNormalization((0.5, 2.0))
                                   )
                               ]),
                               sometimes(iaa.PiecewiseAffine(scale=(0.01, 0.05))),
                               # sometimes move parts of the image around

                           ],random_order=True)
            ],random_order=True
        )
    def augment(self,batch=16,scale=2):
        n_step = self._len() // batch
        count = 0
        while True:
            data = self._load_batch(batch)
            images_aug = self.seq.augment_images(data)
            for i in range(len(images_aug)):
                imdata = images_aug[i]
                misc.imsave(self._random_name()+".png",imdata)
            if count > scale*n_step:
                break
            count += 1


if __name__ == "__main__":
    odir = r"C:\Users\Administrator\Desktop\SSD\train1\NG"
    ydir = r'C:\Users\Administrator\Desktop\SSD\result\ng'
    iau = IAU(odir,ydir)
    iau.augment()


